#include <stdio.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/poll.h>
#include <sys/time.h>
#include <errno.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <list>


// [[ex07aa]]
// DEBUG descr The API of a poll-based library.
class library 
{
public:
	library()
	{ }
	virtual ~library()
	{ }

	// get the number of file descriptors that need to
	// be polled
	virtual int get_n_desc()
	{ 
		return 0;
	}
	// fill out the information about the file
	// descriptors to be polled
	virtual void fill_desc(pollfd *fds, int fdcount)
	{ }
	// get the information about the time of the next
	// timeout, now is the current time, for
	// convenience of relative computations; returns
	// true if any, false if no timeout
	virtual bool get_next_timeout(timeval &ts, 
		const timeval &now)
	{
		return false;
	}
	// process any pending file and timeout events
	virtual void process(pollfd *fds, int fdcount, 
		const timeval &now)
	{ }
};
// END

// [[ex07ab]]
// DEBUG descr The event loop for the poll-based libraries.
class event_loop
{
protected:
	struct libhandle {
		// the virtual interface
		library *lib_; 
		// number returned by get_n_desc()
		int ndesc_; 
		// offset in the common descriptor array
		int first_desc_; 
		// as returned by get_next_timeout()
		timeval timeout_; 
		// as returned by get_next_timeout()
		bool has_timeout_; 
	};

	// units to grow the internal array of pollfds
	static const int FDROUND = 16;
public:
	event_loop() :
		nlibs_(0), fds_(0), nfds_(0), handles_(0), 
		nhandles_(0)
	{ }
	~event_loop()
	{ 
		delete [] fds_;
		delete [] handles_;
		// the libraries that weren't forgotten are
		// destroyed
		for (liblist::iterator it = libs_.begin(); 
				it != libs_.end(); ++it)
			delete *it;
	}

	// the caller must be careful not to add the same
	// library twice
	void add_library(library *lib)
	{
		libs_.push_back(lib);
		nlibs_++;
	}
	void forget_library(library *lib)
	{
		for (liblist::iterator it = libs_.begin(); 
				it != libs_.end(); ++it) {
			if (*it == lib) {
				libs_.erase(it);
				nlibs_--;
				break;
			}
		}
		for (int i = 0; i < nhandles_; i++) {
			if (handles_[i].lib_ == lib) {
				handles_[i].lib_ = 0;
				break;
			}
		}
	}
	// one iteration of the loop, returns false on
	// poll error
	bool run_once()
	{
		timeval now, nextto;
		bool hastimeout = false;
		gettimeofday(&now, NULL); // the current time

		if (nlibs_ > nhandles_) {
			delete handles_;
			nhandles_  = nlibs_;
			handles_ = new libhandle[nhandles_];
		}

		// copy the libraries list to the array of
		// handles that will stay stable during this one
		// run
		int i = 0;
		for (liblist::iterator it = libs_.begin(); 
				it != libs_.end(); ++it) {
			handles_[i++].lib_ = *it;
		}

		// total number of fd structures needed
		int nfds = 0; 

		for (i = 0; i < nhandles_; i++) {
			libhandle *lh = &handles_[i];
			library *lib = lh->lib_;
			lh->ndesc_ = lib->get_n_desc();
			lh->first_desc_ = nfds;
			nfds += lh->ndesc_;

			lh->has_timeout_ = lib->get_next_timeout(
				lh->timeout_, now);
			// merge into the common timeout information)
			if (lh->has_timeout_) {
				if (hastimeout) {
					if (lh->timeout_.tv_sec < nextto.tv_sec
					|| lh->timeout_.tv_sec == nextto.tv_sec
						&& lh->timeout_.tv_usec < nextto.tv_usec)
						nextto = lh->timeout_;
				} else {
					hastimeout = true;
					nextto = lh->timeout_;
				}
			}
		}
		
		if (nfds_ < nfds) {
			nfds_ = (nfds/FDROUND + 1)*FDROUND;
			delete [] fds_;
			fds_ = new pollfd[nfds_];
		}

		// now collect the information about the descriptors
		memset(fds_, 0, sizeof(pollfd) * nfds_);
		for (i = 0; i < nhandles_; i++) {
			libhandle *lh = &handles_[i];
			lh->lib_->fill_desc(fds_ + lh->first_desc_, 
				lh->ndesc_);
		}

		// calculate the amount of time to sleep
		
		// time in milliseconds, -1 is "infinity"
		int tosleep = -1; 

		if (hastimeout) {
			tosleep = (nextto.tv_sec - now.tv_sec) * 1000 
				+ (nextto.tv_usec - now.tv_usec) / 1000;
			if (tosleep < 0) {
				if (nextto.tv_sec <= now.tv_sec)
					// the next requested time is in the past
					tosleep = 0; 
				else
					// overflow happened, limit to 1 day
					tosleep = 24*3600*1000; 
			} else {
				// round the timeout up, or it will be
				// busy-looping for 1 ms when the time comes
				int diff = nextto.tv_usec - now.tv_usec;
				if (diff < 0)
					diff += 1000*1000;
				if (diff % 1000)
					tosleep++;
			}
		}
	
		if (poll(fds_, nfds, tosleep) < 0) {
			if (errno == EINTR || errno == EAGAIN)
				return true;
			else
				return false;
		}

		// now handle the results
		gettimeofday(&now, NULL); // the current time
		for (i = 0; i < nhandles_; i++) {
			libhandle *lh = &handles_[i];
			// the libraries may get deleted, then skip them
			if (lh->lib_  == 0) 
				continue;
			lh->lib_->process(fds_ + lh->first_desc_, 
				lh->ndesc_, now);
		}
		return true;
	}
	// an (almost) endless loop, returns false on poll
	// error
	bool run_loop()
	{
		while(nlibs_ > 0) {
			if (!run_once())
				return false;
		}
		return true;
	}

protected:
	typedef std::list<library *> liblist;
	liblist libs_; 
	int nlibs_; // size of libs_
	pollfd *fds_;
	int nfds_; // size of fds_ 
	libhandle *handles_;
	int nhandles_; // size of handles_
};
// END

// for collecting the data to be written to the socket
// [[ex07ac]]
// DEBUG descr The simple chained buffer.
class buffer 
{
protected:
	enum {
		CHUNK_SIZE = 512
	};
	class chunk
	{
	public:
		chunk() :
			wpos_(0), rpos_(0), data_(new char[CHUNK_SIZE])
		{ }
		~chunk() 
		{
			delete [] data_;
		}
		int size()
		{
			return wpos_ - rpos_;
		}
		int freesize()
		{
			return CHUNK_SIZE - wpos_;
		}
		void append(char *data, int n)
		{
			memcpy(data_ + wpos_, data, n);
			wpos_ += n;
		}
		char *getdata()
		{
			return data_ + rpos_;
		}
		// returns the size of data left in the chunk
		int consumed(int n)
		{
			rpos_ += n;
			return wpos_ - rpos_;
		}

	protected:
		int wpos_; // write position
		int rpos_; // read position
		char *data_; // the actual data buffer
	};

public:
	buffer() :
		size_(0)
	{ }
	~buffer()
	{
		for (chunklist::iterator it = chunks_.begin(); 
				it != chunks_.end(); ++it) {
			delete *it;
		}
	}

	// the total size of data in the buffer
	int size()
	{
		return size_;
	}

	// append a block of data of n bytes
	void append(char *data, int n)
	{
		size_ += n;

		if (chunks_.empty()) 
			chunks_.push_back(new chunk);

		// if there is space left in the last chunk, fill it
		chunk *lastch = chunks_.back();
		int left = lastch->freesize();
		if (left > n)
			left = n;
		lastch->append(data, left);
		n -= left;

		// then add as many chunks as needed
		while(n > 0) {
			lastch = new chunk;
			chunks_.push_back(lastch);
			left = lastch->freesize();
			if (left > n)
				left = n;
			lastch->append(data, left);
			n -= left;
		}
	}

	// the following returns the values convenient for
	// passing the first chunk to the write() call;
	// buffer must not be empty
	char *contigdata()
	{
		return chunks_.front()->getdata();
	}
	int contigsize()
	{
		return chunks_.front()->size();
	}
	// indication that an amount of data was consumed
	// by the write() call
	void consumed(int n)
	{
		if (n < 0)
			// it was an error indication, not data size
			return; 
		chunk *ch = chunks_.front();
		if (ch->consumed(n) <= 0) { 
			// front chunk emptied out, free it
			chunks_.pop_front();
			delete ch;
		}
		size_ -= n;
	}

protected:
	typedef std::list<chunk *> chunklist;
	chunklist chunks_;
	int size_;
};
// END

// [[ex07ad]]
// DEBUG descr A simple poll-based network server.
class server_echo : public library
{
protected:
	enum {
		// if the write buffer has this many bytes, stop
		// reading that socket
		BUFFER_LIMIT = 10 // 10240 
	};

	class client {
	public:
		client(int fd = -1, const timeval *now = 0) :
			fd_(fd), reof_(false)
		{ 
			if (now)
				lastused_ = *now;
			else
				// the current time
				gettimeofday(&lastused_, NULL); 
		}
		~client()
		{
			if (fd_ >= 0)
				close(fd_);
		}
		buffer wbf_; // the write buffer
		int fd_; // the socket file descriptor
		bool reof_; // the reading side received EOF
		timeval lastused_; // when the socket was last read
	};

	typedef std::list<client *> clist;

private:
	server_echo();

public:
	// the library will be listening on this port
	server_echo(int port, event_loop *evloop, 
			int naccepts = -1) :
		port_(port), srvsock_(-1), nclients_(0),
		evloop_(evloop), naccepts_(naccepts)
	{
		srvsock_ = socket(PF_INET, SOCK_STREAM, 0);
		// bind() and listen() are quick,
		// for any following calls set the NONBLOCK flag
		if (srvsock_ >= 0) { 
			int one = 1;
			setsockopt(srvsock_, SOL_SOCKET, SO_REUSEADDR, 
				(char *)&one, sizeof(one));
			sockaddr_in addr;
			addr.sin_family = AF_INET;
			addr.sin_port = htons(port_);
			addr.sin_addr.s_addr = INADDR_ANY;
			if (bind(srvsock_, 
					(sockaddr *)&addr, sizeof addr) < 0
			|| listen(srvsock_, 5) < 0
			|| fcntl(srvsock_, F_SETFL, O_NONBLOCK) < 0) {
				close(srvsock_);
				srvsock_ = -1;
			}
		}
	}

	~server_echo()
	{
		if (srvsock_ >= 0)
			close(srvsock_);
		for (clist::iterator it = clients_.begin(); 
				it != clients_.end(); ++it)
			delete *it;
	}

	virtual int get_n_desc()
	{ 
		return nclients_ + (srvsock_ >= 0 ? 1 : 0);
	}
	// fill out the information about the file
	// descriptors to be polled
	virtual void fill_desc(pollfd *fds, int fdcount)
	{ 
		fprintf(stderr, "filling %d descriptors\n", 
			fdcount);
		if (srvsock_ >= 0) {
			fprintf(stderr, "fd %d is server socket\n", 
				srvsock_);
			fds->fd = srvsock_;
			fds->events = (POLLIN | POLLERR);
			fds++;
		}
		for (clist::iterator it = clients_.begin(); 
				it != clients_.end(); ++it, ++fds) {
			client *cl = *it;
			fds->fd = cl->fd_;
			int bufsz = cl->wbf_.size();
			if (bufsz > 0) {
				fprintf(stderr, "fd %d needs writing\n", 
					cl->fd_);
				fds->events |= POLLOUT;
			}
			if (bufsz < BUFFER_LIMIT && !cl->reof_) {
				fprintf(stderr, "fd %d needs reading\n", 
					cl->fd_);
				fds->events |= POLLIN;
			}
			// Linux always checks for POLLERR, but add it
			// explicitly anyway
			fds->events |= POLLERR;
		}
	}
	// get the information about the time of the next
	// timeout, now is the current time, for
	// convenience of relative computations; returns
	// true if any, false if no timeout
	virtual bool get_next_timeout(timeval &ts, 
		const timeval &now)
	{
		bool needto = false;

		// first find the earliest last time some socket
		// was read
		for (clist::iterator it = clients_.begin(); 
				it != clients_.end(); ++it) {
			client *cl = *it;
			// a client that is not closed and not
			// stalling the writes is suitable for the
			// timeout processing
			if ( !cl->reof_ && cl->wbf_.size() == 0 ) {
				// this client needs a timeout
				if (needto) {
					if (cl->lastused_.tv_sec < ts.tv_sec
					|| cl->lastused_.tv_sec == ts.tv_sec
						&& cl->lastused_.tv_usec < ts.tv_usec)
						ts = cl->lastused_;
				} else {
					needto = true;
					ts = cl->lastused_;
				}
			}
		}
		if (needto) {
			// then wait for the moment 5 seconds from
			// that time
			ts.tv_sec += 5; 
			timeval diff;
			diff.tv_sec = ts.tv_sec - now.tv_sec;
			diff.tv_usec = ts.tv_usec - now.tv_usec;
			if (diff.tv_usec < 0) {
				diff.tv_sec--;
				diff.tv_usec += (1000*1000);
			}
			fprintf(stderr, "needs timeout in %d.%06d\n", 
				diff.tv_sec, diff.tv_usec);
		}
		return needto;
	}
	// process any pending file and timeout events
	virtual void process(pollfd *fds, int fdcount, 
		const timeval &now)
	{ 
		fprintf(stderr, "process\n");
		int newfd = -1;
		if (srvsock_ >= 0) {
			if (fds->revents & (POLLERR|POLLNVAL)) {
				close(srvsock_);
				srvsock_ = -1;
			} else if (fds->revents & POLLIN) {
				sockaddr_in addr;
				socklen_t alen = sizeof(addr);
				newfd = accept(srvsock_, 
					(sockaddr *)&addr, &alen);
				if (naccepts_ >= 0) {
					if (--naccepts_ <= 0) {
						close(newfd);
						evloop_->forget_library(this);
						delete this;
						return;
					}
				}
				if (newfd >= 0) {
					// set the non-blocking mode right away
					if (fcntl(newfd, F_SETFL, O_NONBLOCK) < 0) {
						close(newfd);
						newfd = -1;
					}
				}
				fprintf(stderr, 
					"got a new connection at fd %d\n", newfd);
			}
			fds++;
		}
		for (clist::iterator it = clients_.begin(); 
				it != clients_.end(); ++fds) {
			client *cl = *it;
			int bufsz = cl->wbf_.size();
			if (fds->revents & POLLERR) {
				close(cl->fd_);
				cl->fd_ = -1;
			}
			if (cl->fd_ >= 0 && bufsz > 0 
			&& (fds->revents & POLLOUT)) {
				int res = write(cl->fd_, cl->wbf_.contigdata(), 
					cl->wbf_.contigsize());
				cl->wbf_.consumed(res);
				if (res > 0)
					cl->lastused_ = now;
				fprintf(stderr, "wrote %d bytes to fd %d\n", 
					res, cl->fd_);
				if (res < 0
				&& errno != EINTR && errno != EWOULDBLOCK 
				&& errno != EAGAIN) {
					close(cl->fd_);
					cl->fd_ = -1;
				}
			} 
			if (cl->fd_ >= 0 && bufsz < BUFFER_LIMIT 
			&& (fds->revents & POLLIN)) {
				char buf[1024];
				int res = read(cl->fd_, buf, sizeof(buf));
				fprintf(stderr, "read %d bytes from fd %d\n", 
					res, cl->fd_);
				if (res > 0) {
					cl->lastused_ = now;
					cl->wbf_.append(buf, res);
					fprintf(stderr, 
						"now the write buffer is %d bytes\n", 
						cl->wbf_.size());
				} else if (res == 0) {
					// got an EOF, no more reading
					cl->reof_ = true; 
					fprintf(stderr, 
						"got eof at fd %d, %d bytes pending\n", 
						cl->fd_, cl->wbf_.size());
				} else if (res < 0
				&& errno != EINTR && errno != EWOULDBLOCK 
				&& errno != EAGAIN) {
					close(cl->fd_);
					cl->fd_ = -1;
				}
			}

			if (cl->reof_ && cl->wbf_.size() == 0) {
				fprintf(stderr, "done at fd %d, closing\n", 
					cl->fd_);
				// reading completed, 
				// all the pending output flushed
				close(cl->fd_);
				cl->fd_ = -1;
			}

			if (cl->fd_ >= 0) {
				if ( !cl->reof_ && cl->wbf_.size() == 0
				&& ( (now.tv_sec - cl->lastused_.tv_sec) 
							* (1000*1000)
						+ now.tv_usec - cl->lastused_.tv_usec 
					> (5*1000*1000) )) 
				{
					cl->wbf_.append("\n?\n", 3);
				}
				++it;
			} else {
				fprintf(stderr, "dropping a client\n");
				delete cl;
				clients_.erase(it++);
				--nclients_;
			}
		}
		// if a new connection was accepted, create a
		// structure for it
		if (newfd >= 0) {
			fprintf(stderr, 
				"adding the new client for fd %d\n", newfd);
			clients_.push_back(new client(newfd, &now));
			++nclients_;
		}
	}

protected:
	int port_;
	int srvsock_;
	clist clients_;
	int nclients_;
	event_loop *evloop_;
	int naccepts_;
};
// END

// [[ex07ae]]
// DEBUG descr The main program invoking the network server in an event loop.
int main()
{
	event_loop el;
	el.add_library(new server_echo(12345, &el, 3));
	if (!el.run_loop())
		perror("poll error");
	return 0;
}
// END
